package com.ruoyi.common.datasource;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import javax.sql.DataSource;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.util.CollectionUtils;

/**
 * 动态数据源
 * 
 * @author ruoyi
 */
public class DynamicDataSource extends AbstractRoutingDataSource
{
    private final Map<Object, Object> targetDataSources = new ConcurrentHashMap<>();

    public DynamicDataSource(DataSource defaultTargetDataSource, Map<Object, Object> targetDataSources)
    {
        super.setDefaultTargetDataSource(defaultTargetDataSource);
        if (!CollectionUtils.isEmpty(targetDataSources))
        {
            this.targetDataSources.putAll(targetDataSources);
            super.setTargetDataSources(targetDataSources);
        }
        super.afterPropertiesSet();
    }

    @Override
    protected Object determineCurrentLookupKey()
    {
        return DynamicDataSourceContextHolder.getDataSourceType();
    }

    /**
     * 添加数据源
     * 
     * @param key 数据源标识
     * @param dataSource 数据源
     */
    public void addDataSource(Object key, DataSource dataSource)
    {
        this.targetDataSources.put(key, dataSource);
        super.setTargetDataSources(this.targetDataSources);
        super.afterPropertiesSet();
    }

    /**
     * 删除数据源
     * 
     * @param key 数据源标识
     */
    public void removeDataSource(Object key)
    {
        this.targetDataSources.remove(key);
        super.setTargetDataSources(this.targetDataSources);
        super.afterPropertiesSet();
    }

    /**
     * 清空数据源
     */
    public void clearDataSource()
    {
        this.targetDataSources.clear();
        super.setTargetDataSources(this.targetDataSources);
        super.afterPropertiesSet();
    }

    /**
     * 获取数据源
     * 
     * @param key 数据源标识
     * @return 数据源
     */
    public DataSource getDataSource(Object key)
    {
        return (DataSource) this.targetDataSources.get(key);
    }

    /**
     * 获取所有数据源
     * 
     * @return 数据源集合
     */
    public Map<Object, Object> getAllDataSources()
    {
        return new ConcurrentHashMap<>(this.targetDataSources);
    }
}